package club.janna.springboot.study.interceptor;

import club.janna.springboot.study.annotation.JannaLimiter;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

@Slf4j
public class JannaLimiterInterceptor extends HandlerInterceptorAdapter {

    private final Map<String, RateLimiter> map = new ConcurrentHashMap<>();

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if(handler instanceof HandlerMethod) {
            HandlerMethod method = (HandlerMethod) handler;
            // 获取到限流的注解
            JannaLimiter jannaLimiter = method.getMethodAnnotation(JannaLimiter.class);
            if(jannaLimiter == null)
                return true; // 没有加限流的注解，通过
            // 获取请求的uri
            String uri = request.getRequestURI();
            RateLimiter limiter;
            if(!map.containsKey(uri))
                synchronized (map) {
                    if(!map.containsKey(uri)) { // 不存在rateLimiter，创建一个
                        log.info("创建RateLimiter，key -> {}, permitsPerSecond -> {}", uri, jannaLimiter.value());
                        map.putIfAbsent(uri, RateLimiter.create(jannaLimiter.value()));
                    }
                }
            limiter = map.get(uri);
            if(!limiter.tryAcquire(jannaLimiter.timeout(), TimeUnit.MILLISECONDS))
                throw new RuntimeException(jannaLimiter.msg()); // 限流
        }
        return true;
    }
}
